import flask
from celery import Celery


class FlaskCelery(Celery):
    def __init__(self, app=None, *args, **kwargs):
        super(FlaskCelery, self).__init__(*args, **kwargs)
        self.patch_task()
        if app is not None:
            self.init_app(app)

    def init_app(self, app):
        self.app = app
        self.config_from_object(app.config['CELERY'])

    def patch_task(self):
        task_base = self.Task
        _celery = self

        class ContextTask(task_base):
            def __call__(self, *args, **kwargs):
                if flask.has_app_context():
                    return task_base.__call__(self, *args, **kwargs)
                else:
                    with _celery.app.app_context():
                        return task_base.__call__(self, *args, **kwargs)

        setattr(ContextTask, 'abstract', True)
        setattr(self, 'Task', ContextTask)


class _CeleryState(object):
    def __init__(self, celery, app):
        self.celery = celery
        self.app = app
